package drds.binlog.database_driver;

import drds.binlog.common.Authors;
import drds.binlog.database_driver.packets.HeaderPacket;
import drds.binlog.database_driver.packets.client.QueryCommandPacket;
import drds.binlog.database_driver.packets.server.*;
import drds.binlog.database_driver.socket.SocketChannel;
import drds.binlog.database_driver.utils.Packets;
import drds.common.Author;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 默认输出的数据编码为UTF-8，如有需要请正确转码
 */
@Author(name = Authors.LI_YANG)
public class QueryExecutor
{

    private SocketChannel socketChannel;

    public QueryExecutor(Connector connector) throws IOException
    {
        if (!connector.isConnected())
        {
            throw new IOException("should execute connector.connect() first");
        }

        this.socketChannel = connector.getSocketChannel();
    }

    public QueryExecutor(SocketChannel socketChannel)
    {
        this.socketChannel = socketChannel;
    }

    /**
     * (Result Set Header Packet) the number of columns <br>
     * (Field Packets) column descriptors <br>
     * (EOF Packet) marker: end of Field Packets <br>
     * (Row Data Packets) row contents <br>
     * (EOF Packet) marker: end of Data Packets
     *
     * @param queryString
     * @return
     * @throws IOException
     */
    public ResultSetPacket query(String queryString) throws IOException
    {
        QueryCommandPacket queryCommandPacket = new QueryCommandPacket();
        queryCommandPacket.setQueryString(queryString);
        byte[] bodyBytes = queryCommandPacket.toBytes();
        Packets.writeBody(socketChannel, bodyBytes);
        byte[] body = readNextPacketBody();

        if (body[0] < 0)
        {
            ErrorPacket errorPacket = new ErrorPacket();
            errorPacket.fromBytes(body);
            throw new IOException(errorPacket + "\n with command: " + queryString);
        }

        ResultSetHeaderPacket resultSetHeaderPacket = new ResultSetHeaderPacket();
        resultSetHeaderPacket.fromBytes(body);

        List<FieldPacket> fieldPacketList = new ArrayList<FieldPacket>();
        for (int i = 0; i < resultSetHeaderPacket.getColumnCount(); i++)
        {
            FieldPacket fieldPacket = new FieldPacket();
            fieldPacket.fromBytes(readNextPacketBody());
            fieldPacketList.add(fieldPacket);
        }

        readEofPacket();

        List<RowDataPacket> rowDataPacketList = new ArrayList<RowDataPacket>();
        while (true)
        {
            body = readNextPacketBody();
            if (body[0] == -2)
            {
                break;
            }
            RowDataPacket rowDataPacket = new RowDataPacket();
            rowDataPacket.fromBytes(body);
            rowDataPacketList.add(rowDataPacket);
        }

        ResultSetPacket resultSetPacket = new ResultSetPacket();
        resultSetPacket.getFieldDescriptorList().addAll(fieldPacketList);
        for (RowDataPacket rowDataPacket : rowDataPacketList)
        {
            resultSetPacket.getFieldValueList().addAll(rowDataPacket.getColumnNameList());
        }
        resultSetPacket.setSourceAddress(socketChannel.getRemoteSocketAddress());

        return resultSetPacket;
    }

    public List<ResultSetPacket> querys(String queryString) throws IOException
    {
        QueryCommandPacket queryCommandPacket = new QueryCommandPacket();
        queryCommandPacket.setQueryString(queryString);
        byte[] bodyBytes = queryCommandPacket.toBytes();
        Packets.writeBody(socketChannel, bodyBytes);
        List<ResultSetPacket> resultSets = new ArrayList<ResultSetPacket>();
        boolean moreResult = true;
        while (moreResult)
        {
            byte[] body = readNextPacketBody();
            if (body[0] < 0)
            {
                ErrorPacket packet = new ErrorPacket();
                packet.fromBytes(body);
                throw new IOException(packet + "\n with command: " + queryString);
            }

            ResultSetHeaderPacket resultSetHeaderPacket = new ResultSetHeaderPacket();
            resultSetHeaderPacket.fromBytes(body);

            List<FieldPacket> fieldPacketList = new ArrayList<FieldPacket>();
            for (int i = 0; i < resultSetHeaderPacket.getColumnCount(); i++)
            {
                FieldPacket fieldPacket = new FieldPacket();
                fieldPacket.fromBytes(readNextPacketBody());
                fieldPacketList.add(fieldPacket);
            }

            moreResult = readEofPacket();

            List<RowDataPacket> rowDataPacketList = new ArrayList<RowDataPacket>();
            while (true)
            {
                body = readNextPacketBody();
                if (body[0] == -2)
                {
                    break;
                }
                RowDataPacket rowDataPacket = new RowDataPacket();
                rowDataPacket.fromBytes(body);
                rowDataPacketList.add(rowDataPacket);
            }

            ResultSetPacket resultSetPacket = new ResultSetPacket();
            resultSetPacket.getFieldDescriptorList().addAll(fieldPacketList);
            for (RowDataPacket rowDataPacket : rowDataPacketList)
            {
                resultSetPacket.getFieldValueList().addAll(rowDataPacket.getColumnNameList());
            }
            resultSetPacket.setSourceAddress(socketChannel.getRemoteSocketAddress());
            resultSets.add(resultSetPacket);
        }

        return resultSets;
    }

    private boolean readEofPacket() throws IOException
    {
        byte[] bytes = readNextPacketBody();
        EofPacket eofPacket = new EofPacket();
        eofPacket.fromBytes(bytes);
        if (bytes[0] != -2)
        {
            throw new IOException("EOF Packet is expected, but packet with field_count=" + bytes[0] + " is found.");
        }

        return (eofPacket.statusFlag & 0x0008) != 0;
    }

    protected byte[] readNextPacketBody() throws IOException
    {
        HeaderPacket headerPacket = Packets.readHeader(socketChannel, 4);
        return Packets.readBytes(socketChannel, headerPacket.getPacketBodyLength());
    }
}
